#!/usr/bin/env python3
"""
将json2fsm生成的DOT文件转换为TLS协议的输入输出状态机
基于RFC 8446 (TLS 1.3) 标准定义输入输出消息类型
"""

import re
import json
import argparse
import os
from typing import Dict, List, Tuple, Set
from collections import defaultdict


class TLSMessageMapper:
    """TLS消息类型映射器，基于RFC 8446标准"""
    
    def __init__(self):
        # TLS 1.3 Handshake消息类型 (RFC 8446 Section 4)
        self.tls_messages = {
            # 客户端消息
            'SSL3_MT_CLIENT_HELLO': 'ClientHello',
            'SSL3_MT_CERTIFICATE': 'Certificate', 
            'SSL3_MT_CERTIFICATE_VERIFY': 'CertificateVerify',
            'SSL3_MT_FINISHED': 'Finished',
            'SSL3_MT_END_OF_EARLY_DATA': 'EndOfEarlyData',
            'SSL3_MT_KEY_UPDATE': 'KeyUpdate',
            'SSL3_MT_NEXT_PROTO': 'NextProtocol',
            'SSL3_MT_CLIENT_KEY_EXCHANGE': 'ClientKeyExchange',
            'SSL3_MT_CHANGE_CIPHER_SPEC': 'ChangeCipherSpec',
            
            # 服务器消息
            'SSL3_MT_SERVER_HELLO': 'ServerHello',
            'SSL3_MT_ENCRYPTED_EXTENSIONS': 'EncryptedExtensions',
            'SSL3_MT_CERTIFICATE_REQUEST': 'CertificateRequest',
            'SSL3_MT_COMPRESSED_CERTIFICATE': 'CompressedCertificate',
            'SSL3_MT_NEWSESSION_TICKET': 'NewSessionTicket',
            'SSL3_MT_SERVER_KEY_EXCHANGE': 'ServerKeyExchange',
            'SSL3_MT_SERVER_DONE': 'ServerHelloDone',
            'SSL3_MT_CERTIFICATE_STATUS': 'CertificateStatus',
            'SSL3_MT_HELLO_REQUEST': 'HelloRequest',
            
            # DTLS消息
            'DTLS1_MT_HELLO_VERIFY_REQUEST': 'HelloVerifyRequest',
        }
        
        # 输入输出方向映射
        self.direction_mapping = {
            # 客户端写 (发送)
            'TLS_ST_CW_': 'OUT',
            'TLS_ST_CR_': 'IN',
            # 服务器读 (接收)
            'TLS_ST_SR_': 'IN', 
            'TLS_ST_SW_': 'OUT',
            # DTLS
            'DTLS_ST_': 'IN',
        }
    
    def map_message_type(self, condition: str) -> str:
        """将条件中的消息类型映射到标准TLS消息"""
        for ssl_msg, tls_msg in self.tls_messages.items():
            if ssl_msg in condition:
                return tls_msg
        return 'Unknown'
    
    def get_direction(self, state: str) -> str:
        """根据状态名称确定输入输出方向"""
        for prefix, direction in self.direction_mapping.items():
            if state.startswith(prefix):
                return direction
        return 'UNKNOWN'


class DOTParser:
    """解析DOT文件"""
    
    def __init__(self):
        self.states: Set[str] = set()
        self.transitions: List[Tuple[str, str, str]] = []
    
    def parse_dot_file(self, dot_file_path: str):
        """解析DOT文件，提取状态和转换"""
        with open(dot_file_path, 'r', encoding='utf-8') as f:
            content = f.read()
        
        # 提取所有状态
        state_pattern = r'([A-Z_]+)\s*->\s*([A-Z_]+)'
        matches = re.findall(state_pattern, content)
        
        for start_state, end_state in matches:
            self.states.add(start_state)
            self.states.add(end_state)
        
        # 提取转换条件
        transition_pattern = r'([A-Z_]+)\s*->\s*([A-Z_]+)\s*\[label="([^"]+)"\]'
        matches = re.findall(transition_pattern, content)
        
        for start_state, end_state, condition in matches:
            self.transitions.append((start_state, end_state, condition))
            self.states.add(start_state)
            self.states.add(end_state)
    
    def get_states(self) -> Set[str]:
        return self.states
    
    def get_transitions(self) -> List[Tuple[str, str, str]]:
        return self.transitions


class IOStateMachine:
    """输入输出状态机"""
    
    def __init__(self):
        self.states: Set[str] = set()
        self.io_transitions: Dict[str, List[Tuple[str, str, str]]] = defaultdict(list)
        self.message_mapper = TLSMessageMapper()
    
    def add_state(self, state: str):
        """添加状态"""
        self.states.add(state)
        if state not in self.io_transitions:
            self.io_transitions[state] = []
    
    def add_io_transition(self, from_state: str, to_state: str, input_msg: str, output_msg: str):
        """添加输入输出转换"""
        self.add_state(from_state)
        self.add_state(to_state)
        self.io_transitions[from_state].append((input_msg, output_msg, to_state))
    
    def convert_from_dot_parser(self, dot_parser: DOTParser):
        """从DOT解析器转换状态机"""
        for start_state, end_state, condition in dot_parser.get_transitions():
            # 解析条件，提取消息类型
            input_msg = self.message_mapper.map_message_type(condition)
            output_msg = self.message_mapper.map_message_type(condition)
            
            # 根据状态确定输入输出方向
            start_direction = self.message_mapper.get_direction(start_state)
            end_direction = self.message_mapper.get_direction(end_state)
            
            # 确定实际的输入输出消息
            if start_direction == 'OUT' and end_direction == 'IN':
                # 客户端发送，服务器接收
                actual_input = f"IN:{input_msg}" if input_msg != 'Unknown' else "IN:Message"
                actual_output = f"OUT:{output_msg}" if output_msg != 'Unknown' else "OUT:Message"
            elif start_direction == 'IN' and end_direction == 'OUT':
                # 服务器发送，客户端接收
                actual_input = f"IN:{input_msg}" if input_msg != 'Unknown' else "IN:Message"
                actual_output = f"OUT:{output_msg}" if output_msg != 'Unknown' else "OUT:Message"
            else:
                # 其他情况，保持原样
                actual_input = f"IN:{input_msg}" if input_msg != 'Unknown' else "IN:Message"
                actual_output = f"OUT:{output_msg}" if output_msg != 'Unknown' else "OUT:Message"
            
            self.add_io_transition(start_state, end_state, actual_input, actual_output)
    
    def to_dot(self) -> str:
        """转换为DOT格式"""
        dot_lines = ["digraph IO_FSM {"]
        
        for state in sorted(self.states):
            dot_lines.append(f'    {state} [label="{state}"];')
        
        dot_lines.append("")
        
        for from_state, transitions in self.io_transitions.items():
            for input_msg, output_msg, to_state in transitions:
                label = f"{input_msg}/{output_msg}"
                dot_lines.append(f'    {from_state} -> {to_state} [label="{label}"];')
        
        dot_lines.append("}")
        return "\n".join(dot_lines)
    
    def to_json(self) -> Dict:
        """转换为JSON格式"""
        result = {
            "states": list(self.states),
            "transitions": []
        }
        
        for from_state, transitions in self.io_transitions.items():
            for input_msg, output_msg, to_state in transitions:
                result["transitions"].append({
                    "from_state": from_state,
                    "to_state": to_state,
                    "input": input_msg,
                    "output": output_msg
                })
        
        return result


def main():
    parser = argparse.ArgumentParser(description='将DOT文件转换为TLS输入输出状态机')
    parser.add_argument('-i', '--input', required=True, help='输入DOT文件路径')
    parser.add_argument('-o', '--output', required=True, help='输出目录')
    parser.add_argument('-f', '--format', choices=['dot', 'json', 'both'], default='both', 
                       help='输出格式')
    
    args = parser.parse_args()
    
    # 创建输出目录
    os.makedirs(args.output, exist_ok=True)
    
    # 解析DOT文件
    print(f"解析DOT文件: {args.input}")
    dot_parser = DOTParser()
    dot_parser.parse_dot_file(args.input)
    
    print(f"发现 {len(dot_parser.get_states())} 个状态")
    print(f"发现 {len(dot_parser.get_transitions())} 个转换")
    
    # 转换为输入输出状态机
    print("转换为输入输出状态机...")
    io_fsm = IOStateMachine()
    io_fsm.convert_from_dot_parser(dot_parser)
    
    # 生成输出文件
    base_name = os.path.splitext(os.path.basename(args.input))[0]
    
    if args.format in ['dot', 'both']:
        dot_output = os.path.join(args.output, f"{base_name}_io.dot")
        with open(dot_output, 'w', encoding='utf-8') as f:
            f.write(io_fsm.to_dot())
        print(f"生成DOT文件: {dot_output}")
    
    if args.format in ['json', 'both']:
        json_output = os.path.join(args.output, f"{base_name}_io.json")
        with open(json_output, 'w', encoding='utf-8') as f:
            json.dump(io_fsm.to_json(), f, indent=2, ensure_ascii=False)
        print(f"生成JSON文件: {json_output}")
    
    print("转换完成!")


if __name__ == '__main__':
    main()
